package org.msi.bayesiannetwork;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * <p>Represents the implementation of the algorithm for evaluating expression for the {@link BayesianNetwork}.</p>
 *
 * @author Jakub Narloch
 * @version 1.0
 */
public class RecursiveAlgorithm {

    /**
     * <p>Represents the separator pattern used for dividing the variables from the question from those that are
     * observations.</p>
     */
    private static final String SEPARATOR_PATTERN = "\\|";

    /**
     * <p>Represents the variable negation character.</p>
     */
    private static final String NEGATION = "~";

    /**
     * <p>Represents the pattern for extracting the variables from expression.</p>
     */
    private static final Pattern PATTERN = Pattern.compile("((~\\s*)?(\\w+))");

    /**
     * <p>Represents the instance of {@link BayesianNetwork} class.</p>
     */
    private BayesianNetwork bayesianNetwork;

    /**
     * <p>Creates new instance of {@link RecursiveAlgorithm} class with given bayesian network.</p>
     *
     * @param bayesianNetwork the {@link BayesianNetwork} instance to use
     */
    public RecursiveAlgorithm(BayesianNetwork bayesianNetwork) {

        this.bayesianNetwork = bayesianNetwork;
    }

    /**
     * <p>Evaluates the given expression.</p>
     *
     * @param expression the expression to evaluate
     *
     * @return the probability
     *
     * @throws ExpressionEvaluationException if the error occurs during expression evaluation
     */
    public double evaluateExpression(String expression) throws ExpressionEvaluationException {

        String[] parts;
        List<Variable> variables;
        List<Variable> observations;

        try {
            parts = expression.split(SEPARATOR_PATTERN);

            variables = toList(parts[0]);
            observations = new ArrayList<Variable>();

            if (parts.length == 2) {

                observations = toList(parts[1]);
            } else if (parts.length > 2) {

                throw new ExpressionEvaluationException("The expression contains multiple separators.");
            }

            return evaluate(variables, observations);
        } catch (IllegalArgumentException e) {
            throw new ExpressionEvaluationException("Error occurred when evaluating expression.", e);
        } catch (NodeNotFoundException e) {
            throw new ExpressionEvaluationException("Error occurred when evaluating expression.", e);
        }
    }

    /**
     * <p>Evaluates the probability of given variables.</p>
     *
     * @param variables    the list of variables
     * @param observations the list of observations
     *
     * @return the probability of variables
     *
     * @throws NodeNotFoundException if node for the given variable does not exists
     */
    private double evaluate(List<Variable> variables, List<Variable> observations)
            throws NodeNotFoundException {

        List<Boolean> variableValues;
        List<Double> result = new ArrayList<Double>();
        List<Variable> allVariables = getVariables(variables);
        int indexResult = -1;

        for (int i = 0; i < Math.pow(2, variables.size()); i++) {
            variableValues = Utils.intToBooleanList(i, variables.size());

            result.add(calculate(extend(variables, variableValues, allVariables),
                    extend(variables, variableValues, observations)));

            if (isResult(variables, variableValues)) {
                indexResult = i;
            }
        }

        return normalize(indexResult, result);
    }

    /**
     * <p>Calculates the the probability for the given variables.</p>
     *
     * @param variables    the list of variables
     * @param observations the list of observations
     *
     * @return the probability
     *
     * @throws NodeNotFoundException if node for the given variable does not exists
     */
    private double calculate(List<Variable> variables, List<Variable> observations)
            throws NodeNotFoundException {

        double probability;
        Variable variable;
        Variable var;

        // indicates the algorithm ending condition
        if (variables.size() == 0) {

            return 1.0;
        }

        // retrieves the first variable name
        variable = variables.get(0);
        // removes the retrieved element
        variables.remove(0);

        if (!observations.contains(variable)) {

            var = new Variable(variable.getVariableName(), false);
            probability = probability(var, observations) * calculate(new ArrayList<Variable>(variables), extend(var, observations));
            var = new Variable(variable.getVariableName(), true);
            probability += probability(var, observations) * calculate(new ArrayList<Variable>(variables), extend(var, observations));

            return probability;
        } else {

            return probability(variable, observations) *
                    calculate(new ArrayList<Variable>(variables), observations);
        }
    }

    /**
     * <p>Normalizes the probability for the given nodes.</p>
     *
     * @param resultIndex   represents the result index
     * @param probabilities the variable probabilities
     *
     * @return the normalized probability
     */
    private double normalize(int resultIndex, List<Double> probabilities) {

        double sum = 0D;

        for (Double probability : probabilities) {

            sum += probability;
        }

        return probabilities.get(resultIndex) / sum;
    }

    /**
     * <p>Retrieves the 'parent' variables for the given variable.</p>
     *
     * @param variables the list of variable to use
     *
     * @return the list of 'parent' varibles
     *
     * @throws NodeNotFoundException if the node for the given variable does not exists
     */
    private List<Variable> getVariables(List<Variable> variables) throws NodeNotFoundException {

        List<Variable> allVariables;
        List<Variable> result = new ArrayList<Variable>();
        for (Variable variable : variables) {

            allVariables = bayesianNetwork.getAllVariableParents(variable);

            for (Variable var : allVariables) {
                if (!result.contains(var)) {
                    result.add(var);
                }
            }
        }

        return result;
    }

    /**
     * <p>Returns the probability for the given variable.</p>
     *
     * @param variable     the variable for the probability will be calculated
     * @param observations the list of observations
     *
     * @return the probability
     *
     * @throws NodeNotFoundException if the node for the specified variable does not exists
     */
    private double probability(Variable variable, List<Variable> observations) throws NodeNotFoundException {

        return bayesianNetwork.getVariableProbability(variable, observations);
    }

    /**
     * <p>Extends the list of the observations with the given variable.</p> <p>It returns a shallow copy of the given
     * list.</p>
     *
     * @param variable     the variable to use
     * @param observations the list of observations to extend
     *
     * @return a shallow copy containing all the observations from the given list and extended of the given variable
     */
    private static List<Variable> extend(Variable variable, List<Variable> observations) {

        List<Variable> result = new ArrayList<Variable>(observations);

        if(!result.contains(variable)) {
            result.add(variable);
        }

        return result;
    }

    /**
     * <p>Extends the list of the observations with the given variable.</p> <p>It returns a shallow copy of the given
     * list.</p>
     *
     * @param variables      the list variable to use
     * @param variableValues the list variable values to use
     * @param observations   the list of observations to extend
     *
     * @return a shallow copy containing all the observations from the given list and extended of the given variables
     */
    private static List<Variable> extend(List<Variable> variables, List<Boolean> variableValues, List<Variable> observations) {

        List<Variable> result = new ArrayList<Variable>();

        for(Variable observation : observations) {
            if(!variables.contains(observation)) {
                result.add(observation);
            }
        }

        for (int i = 0; i < variables.size() && i < variableValues.size(); i++) {

            result.add(new Variable(variables.get(i).getVariableName(), variableValues.get(i)));
        }

        return result;
    }

    /**
     * <p>Checks if given combination is searched result.</p>
     * @param variables the variable list
     * @param variableValues the variable values
     * @return flag indicating if the given combination is the result
     */
    private boolean isResult(List<Variable> variables, List<Boolean> variableValues) {

        boolean result = true;
        for (int i = 0; i < variables.size() && i < variableValues.size(); i++) {
            if (variables.get(i).isPrimitive() != variableValues.get(i)) {
                result = false;
            }
        }

        return result;
    }

    /**
     * <p>Parses the given string and extracts from it a list of variables.</p>
     *
     * @param variables the string containing the names of variables
     *
     * @return the list of extracted variables
     *
     * @throws ExpressionEvaluationException if any of variables specified in variables does not exists
     */
    private List<Variable> toList(String variables)
            throws ExpressionEvaluationException {

        List<Variable> result = new ArrayList<Variable>();
        Set<String> variableNames = new HashSet<String>(bayesianNetwork.getVariableList());

        Matcher matcher = PATTERN.matcher(variables);
        String variableName;

        while (matcher.find()) {

            variableName = matcher.group(3).trim();
            checkVariableExists(variableName, variableNames);
            result.add(new Variable(variableName, !matcher.group(1).contains(NEGATION)));
        }

        return result;
    }

    /**
     * <p>Checks if variable with given name exists in bayesian network.</p>
     *
     * @param variableName the variable name
     * @param variables    the set containing all variables
     *
     * @throws ExpressionEvaluationException if the variable with given name does not exists
     */
    private static void checkVariableExists(String variableName, Set<String> variables) throws ExpressionEvaluationException {

        if (!variables.contains(variableName)) {

            throw new ExpressionEvaluationException(String.format("The variable with name: %s, does not exists.",
                    variableName));
        }
    }
}
